"""
数据保存模块
负责IQ数据和分析结果的保存，支持.npy和CSV格式
"""

import numpy as np
import pandas as pd
import json
import os
from pathlib import Path
from typing import Dict, Any, Optional, Union
from datetime import datetime
from src.logger import Logger


class DataSaver:
    """数据保存器，支持多种格式的数据保存"""
    
    def __init__(self, config: Dict[str, Any], logger: Optional[Logger] = None):
        """
        初始化数据保存器
        
        Args:
            config: 输出配置
            logger: 日志记录器
        """
        self.config = config
        self.logger = logger or Logger()
        self.data_directory = Path(config.get('data_directory', './data'))
        self.save_raw_npy = config.get('save_raw_npy', True)
        self.save_raw_csv = config.get('save_raw_csv', True)
        self.save_plots = config.get('save_plots', True)
        self.plot_format = config.get('plot_format', 'png')
        
        # 确保数据目录存在
        self.data_directory.mkdir(parents=True, exist_ok=True)
        self.logger.info(f"数据保存目录: {self.data_directory}")
    
    def _generate_timestamp(self) -> str:
        """
        生成时间戳字符串
        
        Returns:
            格式化的时间戳
        """
        return datetime.now().strftime("%Y%m%d_%H%M%S")
    
    def _create_session_directory(self) -> Path:
        """
        创建会话目录
        
        Returns:
            会话目录路径
        """
        session_dir = self.data_directory / f"session_{self._generate_timestamp()}"
        session_dir.mkdir(parents=True, exist_ok=True)
        self.logger.info(f"创建会话目录: {session_dir}")
        return session_dir
    
    def save_iq_data(self, iq_data: np.ndarray, sample_rate: float,
                    center_freq: float, metadata: Optional[Dict[str, Any]] = None,
                    time_offsets: Optional[np.ndarray] = None,
                    timing_metadata: Optional[Dict[str, Any]] = None) -> Optional[Dict[str, str]]:
        """
        保存IQ数据（包含时间信息）
        
        Args:
            iq_data: 复数IQ数据
            sample_rate: 采样率 (Hz)
            center_freq: 中心频率 (Hz)
            metadata: 元数据
            time_offsets: 时间偏移数组
            timing_metadata: 时序元数据
            
        Returns:
            保存的文件路径字典，如果保存失败则返回None
        """
        try:
            # 创建会话目录
            session_dir = self._create_session_directory()
            base_filename = f"iq_data_{self._generate_timestamp()}"
            
            saved_files = {}
            
            # 保存为.npy格式
            if self.save_raw_npy:
                npy_file = session_dir / f"{base_filename}.npy"
                np.save(npy_file, iq_data)
                saved_files['npy'] = str(npy_file)
                self.logger.info(f"IQ数据已保存为.npy格式: {npy_file}")
            
            # 保存为CSV格式
            if self.save_raw_csv:
                csv_file = session_dir / f"{base_filename}.csv"
                
                # 创建DataFrame，包含实部、虚部和时间信息
                if time_offsets is not None:
                    time_column = time_offsets
                else:
                    time_column = np.arange(len(iq_data)) / sample_rate
                
                df = pd.DataFrame({
                    'index': range(len(iq_data)),
                    'real': iq_data.real,
                    'imag': iq_data.imag,
                    'magnitude': np.abs(iq_data),
                    'phase': np.angle(iq_data),
                    'time_s': time_column,
                    'time_ms': time_column * 1000
                })
                
                df.to_csv(csv_file, index=False)
                saved_files['csv'] = str(csv_file)
                self.logger.info(f"IQ数据已保存为CSV格式: {csv_file}")
            
            # 保存元数据
            if metadata or timing_metadata:
                metadata_file = session_dir / f"{base_filename}_metadata.json"
                
                # 创建基础元数据
                final_metadata = {
                    'data_length': len(iq_data),
                    'sample_rate': sample_rate,
                    'center_frequency': center_freq,
                    'timestamp': datetime.now().isoformat(),
                    'saved_files': saved_files,
                    'has_time_offsets': time_offsets is not None
                }
                
                # 合并用户提供的元数据
                if metadata:
                    final_metadata.update(metadata)
                
                # 添加时序元数据
                if timing_metadata:
                    final_metadata['timing'] = timing_metadata
                    self.logger.info("包含时序元数据")
                
                with open(metadata_file, 'w', encoding='utf-8') as f:
                    json.dump(final_metadata, f, indent=2, ensure_ascii=False)
                
                saved_files['metadata'] = str(metadata_file)
                self.logger.info(f"元数据已保存: {metadata_file}")
            
            return saved_files
            
        except Exception as e:
            self.logger.error(f"保存IQ数据失败: {e}")
            return None
    
    def save_stft_result(self, stft_result: Dict[str, Any], 
                        metadata: Optional[Dict[str, Any]] = None) -> Optional[Dict[str, str]]:
        """
        保存STFT分析结果
        
        Args:
            stft_result: STFT分析结果
            metadata: 元数据
            
        Returns:
            保存的文件路径字典，如果保存失败则返回None
        """
        try:
            # 创建会话目录
            session_dir = self._create_session_directory()
            base_filename = f"stft_result_{self._generate_timestamp()}"
            
            saved_files = {}
            
            # 保存频率和时间轴数据
            if 'f_stft' in stft_result and 't_stft' in stft_result:
                # 保存频率和时间轴数据
                axes_data = {
                    'frequency_stft': stft_result['f_stft'],
                    'time_stft': stft_result['t_stft'],
                    'frequency_offset': stft_result['f_offset'],
                    'frequency_absolute': stft_result['f_abs']
                }
                
                axes_file = session_dir / f"{base_filename}_axes.npy"
                np.save(axes_file, axes_data)
                saved_files['axes_npy'] = str(axes_file)
                
                # 也保存为CSV格式
                axes_df = pd.DataFrame(axes_data)
                axes_csv_file = session_dir / f"{base_filename}_axes.csv"
                axes_df.to_csv(axes_csv_file, index=False)
                saved_files['axes_csv'] = str(axes_csv_file)
                
                self.logger.info(f"频率和时间轴数据已保存")
            
            # 保存STFT功率谱数据
            if 'Sxx_dbm' in stft_result:
                power_file = session_dir / f"{base_filename}_power_dbm.npy"
                np.save(power_file, stft_result['Sxx_dbm'])
                saved_files['power_npy'] = str(power_file)
                
                # 保存原始复数谱数据
                if 'Sxx' in stft_result:
                    power_complex_file = session_dir / f"{base_filename}_power_complex.npy"
                    np.save(power_complex_file, stft_result['Sxx'])
                    saved_files['power_complex_npy'] = str(power_complex_file)
                
                self.logger.info(f"STFT功率谱数据已保存")
            
            # 保存STFT参数
            if 'params' in stft_result:
                params_file = session_dir / f"{base_filename}_params.json"
                with open(params_file, 'w', encoding='utf-8') as f:
                    json.dump(stft_result['params'], f, indent=2, ensure_ascii=False)
                saved_files['params'] = str(params_file)
                self.logger.info(f"STFT参数已保存: {params_file}")
            
            # 保存元数据
            if metadata:
                metadata_file = session_dir / f"{base_filename}_metadata.json"
                
                # 添加STFT信息到元数据
                metadata.update({
                    'stft_timestamp': datetime.now().isoformat(),
                    'saved_files': saved_files
                })
                
                with open(metadata_file, 'w', encoding='utf-8') as f:
                    json.dump(metadata, f, indent=2, ensure_ascii=False)
                
                saved_files['metadata'] = str(metadata_file)
                self.logger.info(f"STFT元数据已保存: {metadata_file}")
            
            return saved_files
            
        except Exception as e:
            self.logger.error(f"保存STFT结果失败: {e}")
            return None
    
    def save_analysis_stats(self, stats: Dict[str, Any], 
                           metadata: Optional[Dict[str, Any]] = None) -> Optional[str]:
        """
        保存分析统计信息
        
        Args:
            stats: 统计信息
            metadata: 元数据
            
        Returns:
            保存的文件路径，如果保存失败则返回None
        """
        try:
            # 创建会话目录
            session_dir = self._create_session_directory()
            filename = f"analysis_stats_{self._generate_timestamp()}.json"
            stats_file = session_dir / filename
            
            # 添加时间戳到统计信息
            stats_with_timestamp = {
                'timestamp': datetime.now().isoformat(),
                'statistics': stats
            }
            
            # 添加元数据
            if metadata:
                stats_with_timestamp['metadata'] = metadata
            
            with open(stats_file, 'w', encoding='utf-8') as f:
                json.dump(stats_with_timestamp, f, indent=2, ensure_ascii=False)
            
            self.logger.info(f"分析统计信息已保存: {stats_file}")
            return str(stats_file)
            
        except Exception as e:
            self.logger.error(f"保存分析统计信息失败: {e}")
            return None
    
    def save_plot(self, figure, plot_name: str, session_dir: Optional[Path] = None) -> Optional[str]:
        """
        保存图表
        
        Args:
            figure: matplotlib图形对象
            plot_name: 图表名称
            session_dir: 会话目录，如果为None则创建新目录
            
        Returns:
            保存的文件路径，如果保存失败则返回None
        """
        try:
            if not self.save_plots:
                return None
            
            if session_dir is None:
                session_dir = self._create_session_directory()
            
            filename = f"{plot_name}_{self._generate_timestamp()}.{self.plot_format}"
            plot_file = session_dir / filename
            
            figure.savefig(plot_file, dpi=300, bbox_inches='tight')
            self.logger.info(f"图表已保存: {plot_file}")
            return str(plot_file)
            
        except Exception as e:
            self.logger.error(f"保存图表失败: {e}")
            return None
    
    def save_complete_session(self, iq_data: np.ndarray, stft_result: Dict[str, Any], 
                            stats: Dict[str, Any], sample_rate: float, 
                            center_freq: float, config: Dict[str, Any]) -> Optional[Dict[str, str]]:
        """
        保存完整的会话数据
        
        Args:
            iq_data: IQ数据
            stft_result: STFT分析结果
            stats: 统计信息
            sample_rate: 采样率
            center_freq: 中心频率
            config: 系统配置
            
        Returns:
            保存的文件路径字典，如果保存失败则返回None
        """
        try:
            # 创建会话目录
            session_dir = self._create_session_directory()
            self.logger.info(f"开始保存完整会话数据到: {session_dir}")
            
            saved_files = {}
            
            # 保存配置
            config_file = session_dir / "system_config.json"
            with open(config_file, 'w', encoding='utf-8') as f:
                json.dump(config, f, indent=2, ensure_ascii=False)
            saved_files['config'] = str(config_file)
            
            # 保存IQ数据
            iq_metadata = {
                'session_id': session_dir.name,
                'data_type': 'iq_complex',
                'sample_rate': sample_rate,
                'center_frequency': center_freq
            }
            iq_files = self.save_iq_data(iq_data, sample_rate, center_freq, iq_metadata)
            if iq_files:
                saved_files['iq_data'] = iq_files
            
            # 保存STFT结果
            stft_metadata = {
                'session_id': session_dir.name,
                'data_type': 'stft_analysis',
                'sample_rate': sample_rate,
                'center_frequency': center_freq
            }
            stft_files = self.save_stft_result(stft_result, stft_metadata)
            if stft_files:
                saved_files['stft_result'] = stft_files
            
            # 保存统计信息
            stats_metadata = {
                'session_id': session_dir.name,
                'data_type': 'analysis_statistics'
            }
            stats_file = self.save_analysis_stats(stats, stats_metadata)
            if stats_file:
                saved_files['statistics'] = stats_file
            
            # 创建会话摘要
            summary_file = session_dir / "session_summary.json"
            session_summary = {
                'session_id': session_dir.name,
                'timestamp': datetime.now().isoformat(),
                'data_info': {
                    'iq_data_length': len(iq_data),
                    'sample_rate': sample_rate,
                    'center_frequency': center_freq
                },
                'analysis_info': {
                    'peak_power_dbm': stats.get('peak_power_dbm'),
                    'avg_power_dbm': stats.get('avg_power_dbm'),
                    'freq_range_mhz': tuple(f/1e6 for f in stats.get('freq_offset_range', (0, 0))),
                    'time_range_s': stats.get('time_range', (0, 0))
                },
                'saved_files': saved_files
            }
            
            with open(summary_file, 'w', encoding='utf-8') as f:
                json.dump(session_summary, f, indent=2, ensure_ascii=False)
            
            saved_files['summary'] = str(summary_file)
            
            self.logger.info("=" * 60)
            self.logger.info("完整会话数据保存完成")
            self.logger.info(f"会话目录: {session_dir}")
            self.logger.info(f"保存的文件: {len(saved_files)} 个")
            for file_type, file_path in saved_files.items():
                self.logger.info(f"  {file_type}: {file_path}")
            self.logger.info("=" * 60)
            
            return saved_files
            
        except Exception as e:
            self.logger.error(f"保存完整会话数据失败: {e}")
            return None
    
    def get_data_directory(self) -> Path:
        """
        获取数据目录路径
        
        Returns:
            数据目录路径
        """
        return self.data_directory
    
    def list_saved_sessions(self) -> list:
        """
        列出已保存的会话
        
        Returns:
            会话目录列表
        """
        try:
            session_dirs = [d for d in self.data_directory.iterdir() 
                          if d.is_dir() and d.name.startswith('session_')]
            session_dirs.sort(reverse=True)  # 最新的会话在前
            return session_dirs
        except Exception as e:
            self.logger.error(f"列出会话失败: {e}")
            return []
    
    def save_chirp_matrix(self, chirp_matrix: np.ndarray, processing_info: Dict[str, Any],
                         sample_rate: float, center_freq: float) -> Optional[Dict[str, str]]:
        """
        保存快慢时间二维矩阵数据
        
        Args:
            chirp_matrix: 快慢时间二维矩阵 [慢时间 × 快时间]
            processing_info: 处理信息
            sample_rate: 采样率 (Hz)
            center_freq: 中心频率 (Hz)
            
        Returns:
            保存的文件路径字典，如果保存失败则返回None
        """
        try:
            # 创建会话目录
            session_dir = self._create_session_directory()
            base_filename = f"chirp_matrix_{self._generate_timestamp()}"
            
            saved_files = {}
            
            # 保存为.npy格式
            if self.save_raw_npy:
                npy_file = session_dir / f"{base_filename}.npy"
                np.save(npy_file, chirp_matrix)
                saved_files['npy'] = str(npy_file)
                self.logger.info(f"快慢时间矩阵已保存为.npy格式: {npy_file}")
            
            # 保存为CSV格式（展平处理，包含慢时间和快时间索引）
            if self.save_raw_csv:
                csv_file = session_dir / f"{base_filename}.csv"
                
                # 获取时间轴信息
                time_axes = processing_info.get('time_axes', {})
                slow_time_axis = time_axes.get('slow_time_axis_ms', np.array([]))
                fast_time_axis = time_axes.get('fast_time_axis_us', np.array([]))
                
                # 创建DataFrame
                rows = []
                for slow_idx in range(chirp_matrix.shape[0]):
                    for fast_idx in range(chirp_matrix.shape[1]):
                        row = {
                            'slow_time_index': slow_idx,
                            'fast_time_index': fast_idx,
                            'slow_time_ms': slow_time_axis[slow_idx] if slow_idx < len(slow_time_axis) else None,
                            'fast_time_us': fast_time_axis[fast_idx] if fast_idx < len(fast_time_axis) else None,
                            'real': chirp_matrix[slow_idx, fast_idx].real,
                            'imag': chirp_matrix[slow_idx, fast_idx].imag,
                            'magnitude': np.abs(chirp_matrix[slow_idx, fast_idx]),
                            'phase': np.angle(chirp_matrix[slow_idx, fast_idx])
                        }
                        rows.append(row)
                
                df = pd.DataFrame(rows)
                df.to_csv(csv_file, index=False)
                saved_files['csv'] = str(csv_file)
                self.logger.info(f"快慢时间矩阵已保存为CSV格式: {csv_file}")
            
            # 保存时间轴数据
            if time_axes:
                axes_file = session_dir / f"{base_filename}_axes.npy"
                axes_data = {
                    'slow_time_axis_ms': slow_time_axis,
                    'fast_time_axis_us': fast_time_axis
                }
                np.save(axes_file, axes_data)
                saved_files['axes_npy'] = str(axes_file)
                
                # 也保存为JSON格式
                axes_json_file = session_dir / f"{base_filename}_axes.json"
                axes_json_data = {
                    'slow_time_axis_ms': slow_time_axis.tolist() if slow_time_axis.size > 0 else [],
                    'fast_time_axis_us': fast_time_axis.tolist() if fast_time_axis.size > 0 else [],
                    'metadata': {
                        'matrix_shape': chirp_matrix.shape,
                        'sample_rate_hz': sample_rate,
                        'center_frequency_hz': center_freq
                    }
                }
                with open(axes_json_file, 'w', encoding='utf-8') as f:
                    json.dump(axes_json_data, f, indent=2, ensure_ascii=False)
                saved_files['axes_json'] = str(axes_json_file)
                self.logger.info("时间轴数据已保存")
            
            # 保存处理信息元数据
            metadata_file = session_dir / f"{base_filename}_metadata.json"
            
            # 创建元数据（确保所有numpy数组都转换为列表）
            processing_info_serializable = {}
            for key, value in processing_info.items():
                if key == 'time_axes':
                    time_axes_serializable = {}
                    for sub_key, sub_value in value.items():
                        if hasattr(sub_value, 'tolist'):
                            time_axes_serializable[sub_key] = sub_value.tolist()
                        else:
                            time_axes_serializable[sub_key] = sub_value
                    processing_info_serializable[key] = time_axes_serializable
                elif hasattr(value, 'tolist'):
                    processing_info_serializable[key] = value.tolist()
                else:
                    processing_info_serializable[key] = value
            
            metadata = {
                'data_type': 'fast_slow_time_matrix',
                'matrix_shape': chirp_matrix.shape,
                'num_chirps': chirp_matrix.shape[0],
                'fast_time_samples': chirp_matrix.shape[1],
                'sample_rate': sample_rate,
                'center_frequency': center_freq,
                'timestamp': datetime.now().isoformat(),
                'saved_files': saved_files,
                'processing_info': processing_info_serializable
            }
            
            with open(metadata_file, 'w', encoding='utf-8') as f:
                json.dump(metadata, f, indent=2, ensure_ascii=False)
            
            saved_files['metadata'] = str(metadata_file)
            self.logger.info(f"快慢时间矩阵元数据已保存: {metadata_file}")
            
            return saved_files
            
        except Exception as e:
            self.logger.error(f"保存快慢时间矩阵失败: {e}")
            return None